#!/usr/bin/env python
#
# Open-source benchmarking suite. Similar to SPEC CPU.
# (c) 2008 Thomas Stromberg <thomas%stromberg.org>
#
# $Id: benchmark.py 771 2008-11-26 03:32:22Z thomas $
import urllib
import os
import random
import re
import shutil
import time
import os.path
import subprocess
import datetime
from optparse import OptionParser

# local
import builder
import util
import yaml
import sysinfo
import version

class FreeBench(object):
  def __init__(self):
    self.current_thread = 0
    self.run_tests = []
    self.run_suites = []
    self.num_attempts = None
    self.compile_only = False
    self.sysinfo = sysinfo.SysInfo()
    self.num_threads = self.sysinfo.ProcessorCount()
    self.base_dir = os.getcwd()
    self.data_dir = os.getcwd() + '/data'
    self.testbase_dir = os.getcwd() + '/test'
    self.LoadConfiguration()
    
  def LoadConfiguration(self):
    """Loads configuration data.

    Side Effects:
      Populates self.tests, self.test_suites, self.recipes
    """
    self.tests = yaml.load(open('tests.yml').read())
    self.test_suites = yaml.load(open('testsuites.yml').read())
  
  def Run(self):
    """Generic Run wrapper."""
    queue = set()
    if not self.num_attempts:
      util.msg("NOTE: Running in quick mode. Use freebench -a3 for official results.")
      self.num_attempts = 1
    if self.run_tests:
      name = 'custom'
      queue.update(self.run_tests)
    if self.run_suites:
      name = '+'.join(self.run_suites)
      for suite in self.run_suites:
        queue.update(self.test_suites[suite])

    if not self.compile_only:  
      build = builder.Builder(base_dir=self.base_dir)
      build.FetchPrecompiledBinaries()

  
    util.msg("Running '%s' test suite with %s tests" % (name, len(queue)))
    results = []
    for test_name in queue:
      if test_name not in self.tests:
        util.msg('Could not find a test named %s, skipping' % test_name)
        break
      best_score = 0
      adjusted_score = 0
      test = self.tests[test_name]
      test_results = self.RunTest(test_name)
      scores = [ x[0] for x in test_results if x ]
      if scores:
        if 'score_parse' in test:
          best_score = max(scores)
        else:
          best_score = min(scores)
      else:
        best_score = None
        adjusted_score = None
        
      if best_score and 'base_score' in test:
        if 'score_parse' in test:
          adjusted_score = (float(best_score / test['base_score'])) * 100
        else:
          adjusted_score = (float(test['base_score']) / best_score) * 100
      else:
        adjusted_score = 0
        
      results.append((test_name, adjusted_score, best_score, test_results))
    self.DisplayConfiguration()
    self.DisplayResults(name, results)
    return results

  def DisplayConfiguration(self):
    si = self.sysinfo
    distro = si.DistroName()
    build = builder.Builder()
    kern = si.KernelName()
    print
    if distro == kern:
      print "%% FreeBench %s - %s %s %s %s" % (version.VERSION, si.DistroName(), si.DistroVersion(), si.Architecture(), build.CompilerBanner())
    else:
      print "%% FreeBench %s - %s %s (%s %s %s)" % (version.VERSION, si.DistroName(), si.DistroVersion(), kern, si.Architecture(), si.KernelVersion())
      print "%% %s" % build.CompilerBanner()
    print "%% %s x %sMHz (%s), %sM RAM, %s" % (si.ProcessorCount(), si.ProcessorSpeedInMhz(), si.ProcessorType(), si.RamInMegabytes(), si.FilesystemType())
    print "%% %s %s" % (build.compiler, build.cflags)
    
  def DisplayResults(self, suite_name, results):
    total_score = 0
    total_best_duration = datetime.timedelta()
    total_adjusted_score = 0
    failures = 0
    attempts = 0
    adjusted_scores = []
    test_variations = []
    print '\n%-18.18s| %6.6s | %-9.9s| %-12.12s | %-6.6s| %s' % (suite_name, 'score (adjusted)', 'best raw', 'best durat.', 'vari.', 'raw scores')
    print '-' * 80
    for (test, adjusted_score, best_score, test_data) in results:
      scores = [ x[0] for x in test_data if x and x[0] ]
      durations = [ x[1] for x in test_data if x and x[1] ]
      threads = [ x[2] for x in test_data if x and x[2] ]
      attempts = len(test_data)
      printed_score = 'N/A'
      if best_score:
        total_score += best_score
      else:
        best_score = 0.0
      if adjusted_score:
        adjusted_scores.append(adjusted_score)
        printed_score = '%5.2f' % adjusted_score
      if scores:
        variation = ((max(scores) / min(scores)) * 100) - 100
      else:
        variation = 0
      test_variations.append(variation)
      if durations:
        best_duration = min(durations)
      else:
        best_duration = datetime.timedelta(seconds=0)
      total_best_duration = total_best_duration + best_duration
      print_scores = [ '%2.2f' % x for x in scores ]
      print '%-18.18s| %-6.6s | %-9.3f| %-12.12s |%5.2f%% | %s' % (test, printed_score, best_score, best_duration, variation, ', '.join(print_scores))
    
    if total_score:
      if adjusted_scores: 
        total_adjusted_score = util.avg(adjusted_scores)
      print '-' * 80
      avg_variation = util.avg(test_variations)
      print '%-18.18s| %-6.2f | %-9.9s| %-12.12s |%5.2f%% |' % ('TOTAL', total_adjusted_score, total_score, total_best_duration, avg_variation)
      print '-' * 80

      if total_adjusted_score:
        ref = "MacBook Pro (Intel Core 2 Duo 2.33GHz)"
        util.msg("This host is %2.2fX as fast as a %s" % ((total_adjusted_score / 100), ref))          
        print
        if attempts == 1:
          util.msg("NOTE: Depending on your machine, there may be ~5% variance between test runs.")
          util.msg("      For better results, use -n3 to the best score of 3 runs")

  def _ExpandMacros(self, string):
    for macro in ('DATA_DIR', 'TEST_DIR', 'NUM_THREADS', 'CURRENT_THREAD',
                  'RANDOM_PORT', 'INPUT_FILE', 'INSTALL_DIR'):
      if macro in string:
        string = string.replace(macro, str(getattr(self, macro.lower())))
    return string


  def BackgroundExec(self, raw_command):
    """Execute a command in the background, expanding any macros."""
    command = self._ExpandMacros(raw_command)
    proc = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE,
                            stderr=subprocess.PIPE, close_fds=True)
    util.debug("  > %s: '%s'" % (proc.pid, command))
    return proc

  def RunTest(self, test_name):
    """Wrapper to build and find binaries, reserve ports, handle multiple runs."""
    results = []
    test = self.tests[test_name]

    recipe_name = test_name.split('.')[0]
    build = builder.Builder(base_dir=self.base_dir)
    try: 
      (install_dir, bin_path) = build.BuildRecipe(recipe_name)
    except ValueError, e:
      util.msg("[FAIL] %s" % e)
      return ((None, None, None))
      
    self.install_dir = install_dir
    util.set_library_path(self.install_dir + '/lib')
    self.test_dir = self.testbase_dir + '/' + test_name
    util.mkdir_p(self.test_dir)
    dir = self._ExpandMacros(test.get('chdir', self.test_dir))
    util.debug("Changing to test directory: %s" % dir)
    util.chdir(dir)

    for attempt in range(self.num_attempts):
      self.random_port = int(random.uniform(16000,32000))
      daemon = None
      try:
        daemon = self._PrepareTest(test_name)
      except ValueError, e:
        util.msg('Could not prepare %s: %s (pwd=%s)' % (test_name, e, os.getcwd()))
        results.append(None)
        break
      
      util.debug("Changing to test directory (just in case): %s" % dir)
      util.chdir(dir)
      (score, delta, threads) = self.ExecuteTest(test_name, bin_path, attempt_num=attempt+1)
      # sleep for sanity.
      if score:
        time.sleep(3)
      if daemon:
        exited = os.waitpid(daemon.pid, os.WNOHANG)
        if exited[0]:
          util.msg('daemon appears to not be running: out=%s err=%s' %
              (daemon.stderr.readlines(), daemon.stdout.readlines()))
          results.append((None, delta, threads))
          break
        else:
          os.kill(daemon.pid, 15)
          time.sleep(1)
          try:
            os.kill(daemon.pid, 9)
          except OSError:
            pass
          daemon.wait()
          util.debug('daemon stdout: %s' % daemon.stdout.readlines())
          util.debug('daemon stderr: %s' % daemon.stderr.readlines())
      results.append((score, delta, threads))
    return results
    
  def _PrepareTest(self, test_name):
    """Creates input files and runs any pre-commands or daemons for a test."""
    test = self.tests[test_name]

    # For a more accurate workload, copy our input file once per thread.
    if 'input_file' in test:
      input_file = self._ExpandMacros(test['input_file'])
      num_threads = test.get('threads', self.num_threads)
      for thread in range(num_threads):
        output_file = '%s/%s.%s' % (self.test_dir, os.path.basename(input_file), thread)
        shutil.copyfile(input_file, output_file)

    if 'pre-command' in test:
      code = self.BackgroundExec(test['pre-command'])
      if code:
        raise ValueError("'%s' exited with code %s" % (test['pre-command'], code))

    if 'daemon' in test:
      return self.BackgroundExec(test['daemon'])
  
  def ExecuteTest(self, test_name, bin_path, attempt_num=None):
    """Executes a test against a particular binary file. Handles multi-core runs."""
    test = self.tests[test_name]
    procs = []
    failed = False
    num_threads = test.get('threads', self.num_threads)    
    command = re.sub('^\w+', bin_path, test['command'])
    util.msg("  Starting '%s' with %s threads  (run %s of %s)" % (test_name, num_threads, attempt_num, self.num_attempts))
    input_base = self.test_dir + '/' + os.path.basename(test.get('input_file', 'NULL')) + '.'
    start_ts = datetime.datetime.now()
    for thread in range(num_threads):
      self.current_thread = thread
      self.input_file = input_base + str(thread)
      procs.append(self.BackgroundExec(command))
 
    for proc in procs:
      exitcode = proc.wait()  
      if exitcode and exitcode != test.get('exitcode', 0):
        failed = True
        util.msg("'%s' exited with %s" % (command, exitcode))
        util.msg('stdout: %s' % proc.stdout.readlines())
        util.msg('stderr: %s' % proc.stderr.readlines())

    end_ts = datetime.datetime.now()
    delta = end_ts - start_ts
    # We return here because we want to reap our children first.
    if failed:
      return (None, delta, num_threads)
    elif 'score_parse' in test:
      score = 0
      for proc in procs:
        output = proc.stdout.readlines()
        match = re.search(test['score_parse'], ''.join(output), re.MULTILINE)
        if match:
          util.debug("Found score match: %s" % match.groups())
          score = score + float(match.groups()[0])
        else:
          util.msg('Failed to match %s in output: %s' % (test['score_parse'], output))
          return None
    else:
      score = util.delta_seconds(delta) / num_threads
    return (score, delta, num_threads)


if __name__ == "__main__": 
  print "*  Freebench %s - $Rev$" % version.VERSION
  parser = OptionParser()
  parser.add_option("-s", "--suite", dest="suite",
                    help="test suite to run")
  parser.add_option("-c", "--compile", dest="compile_only", action="store_true",
                    help="compile all dependencies (do not use packages)")
  parser.add_option("-a", "--number_of_attempts", dest="num_attempts",
                    help="# of attempts to run (default 3)")
  parser.add_option("-v", "--verbose", dest="verbose", action="store_true",
                    help="Turn on verbose mode")
  parser.add_option("-n", "--num_threads", dest="num_threads",
                                      help="number of threads (autodetect)")
  parser.add_option("-t", "--test", dest="test", help="test package to run")
  (options, args) = parser.parse_args()

  fb = FreeBench()
  
  if options.num_attempts:
    fb.num_attempts = int(options.num_attempts)

  if options.compile_only:
    fb.compile_only = True
    
  if options.num_threads:
    fb.num_threads = int(options.num_threads)
    
  if options.verbose:
    VERBOSE = True
  
  if options.suite:
    fb.run_suites = [options.suite]
  elif options.test:
    fb.run_tests = [options.test]
  else:
    fb.run_suites=['standard']
  
  fb.Run()
  
